import time
import os
from prettytable import PrettyTable
from colorama import Fore, Back, Style, init

from stattion_info import station_dict

station_dict_reverse = {x[1]: x[0] for x in station_dict.items()}


class ShowTrainInfo(object):
    def __init__(self):
        self.tb = PrettyTable()
        init(autoreset=True)
        self.tb.field_names = ["车次", "始发站", "终点站", "一等座", "二等座", "软卧", "硬卧", "硬座"]
        self.font_color = [Fore.RED, Fore.GREEN]
    
    @staticmethod
    def _judge_can_order(info):
        if info and info != "无":
            return Fore.BLUE + info + Fore.RESET
        else:
            return "无"
    
    def show(self, trains_info):
        for i, train in enumerate(trains_info):
            self.tb.add_row([self.font_color[i % 2] + train["stationTrainCode"] + Fore.RESET,
                             station_dict_reverse[train["start_station"]],
                             station_dict_reverse[train["end_station"]],
                             self._judge_can_order(train["yidengzuo"]),
                             self._judge_can_order(train["erdengzuo"]),
                             self._judge_can_order(train["ruanwo"]),
                             self._judge_can_order(train["yingwo"]),
                             self._judge_can_order(train["yingzuo"]),
                             ])
        
        print(self.tb)
        self.tb.clear_rows()


if __name__ == '__main__':
    show_12306 = ShowTrainInfo()
    for i in range(10):
        os.system("clear")
        # show_12306.show()
        time.sleep(1)
